

#include "Tracker.hpp"
#include <opencv2/video.hpp>

#include "opencv.hpp"


void Tracker::setLastFrame(cv::Mat frame){
    frame.copyTo(lastFrameGray);
}

void Tracker::setNearestBox(std::vector<cv::Rect> boxes) {
    int index = 0;
    cv::Point p = box.tl() - boxes.at(0).tl();
    int min = p.x * p.x + p.y * p.y;
    for (int i = 1; i < boxes.size(); i++) {
        p = box.tl() - boxes.at(i).tl();
        int d = p.x * p.x + p.y * p.y;
        if (d < min) {
            min = d;
            index = i;
        }
    }
    box = boxes.at(index);
}

void Tracker::detectCorners(cv::Mat &frameGray) {

    // Define tracking region
    cv::Mat trackingRegion = cv::Mat::zeros(frameGray.rows, frameGray.cols, CV_8UC1);
    cv::Point points[1][4];
    points[0][0] = cv::Point(box.tl().x + 0.22 * box.width,
                         box.tl().y + 0.21 * box.height);
    points[0][1] = cv::Point(box.tl().x + 0.78 * box.width,
                         box.tl().y + 0.21 * box.height);
    points[0][2] = cv::Point(box.tl().x + 0.70 * box.width,
                         box.tl().y + 0.65 * box.height);
    points[0][3] = cv::Point(box.tl().x + 0.30 * box.width,
                         box.tl().y + 0.65 * box.height);
    const cv::Point *pts[1] = {points[0]};
    int npts[] = {4};
    cv::fillPoly(trackingRegion, pts, npts, 1, cv::WHITE);

    // Apply corner detection
    cv::goodFeaturesToTrack(frameGray,
                            corners,
                            max_corners,
                            quality_level,
                            min_distance,
                            trackingRegion,
                            3,
                            false,
                            0.04);
}

bool Tracker::trackFace(cv::Mat &frameGray) {

    // Make sure enough corners are available
    if (corners.size() < min_corners) {
        detectCorners(frameGray);
    }

    std::vector<cv::Point2f> corners_1;
    std::vector<cv::Point2f> corners_0;
    std::vector<uchar> cornersFound_1;
    std::vector<uchar> cornersFound_0;
    cv::Mat err;

    // Track face features with Kanade-Lucas-Tomasi (KLT) algorithm
    cv::calcOpticalFlowPyrLK(lastFrameGray, frameGray, corners, corners_1, cornersFound_1, err);

    // Backtrack once to make it more robust
    cv::calcOpticalFlowPyrLK(frameGray, lastFrameGray, corners_1, corners_0, cornersFound_0, err);

    // Exclude no-good corners
    std::vector<cv::Point2f> corners_1v;
    std::vector<cv::Point2f> corners_0v;
    for (size_t j = 0; j < corners.size(); j++) {
        if (cornersFound_1[j] && cornersFound_0[j]
            && cv::norm(corners[j]-corners_0[j]) < 2) {
            corners_0v.push_back(corners_0[j]);
            corners_1v.push_back(corners_1[j]);
        } else {
            std::cout << "Mis!" << std::endl;
        }
    }

    if (corners_1v.size() >= min_corners) {

        // Save updated features
        corners = corners_1v;

        // Estimate affine transform
        cv::Mat transform = cv::estimateRigidTransform(corners_0v, corners_1v, false);

        if (transform.total() > 0) {

            // Update box
            std::vector<cv::Point2f> boxCoords;
            boxCoords.push_back(box.tl());
            boxCoords.push_back(box.br());
            std::vector<cv::Point2f> transformedBoxCoords;

            cv::transform(boxCoords, transformedBoxCoords, transform);
            box = cv::Rect(transformedBoxCoords[0], transformedBoxCoords[1]);

            // Update roi
            std::vector<cv::Point2f> roiCoords;
            roiCoords.push_back(roi.tl());
            roiCoords.push_back(roi.br());
            std::vector<cv::Point2f> transformedRoiCoords;
            cv::transform(roiCoords, transformedRoiCoords, transform);
            roi = cv::Rect(transformedRoiCoords[0], transformedRoiCoords[1]);

            updateMask(frameGray);
        }

    } else {
        std::cout << "Tracking failed! Not enough corners left." << std::endl;
        return false;
    }
    return true;
}

void Tracker::updateROI() {
    this->roi = cv::Rect(cv::Point(box.tl().x + 0.3 * box.width, box.tl().y + 0.1 * box.height),
                     cv::Point(box.tl().x + 0.7 * box.width, box.tl().y + 0.25 * box.height));
}

void Tracker::updateMask(cv::Mat &frameGray) {

    std::cout << "Update mask" << std::endl;

    mask = cv::Mat::zeros(frameGray.size(), frameGray.type());
    cv::rectangle(mask, this->roi, cv::WHITE, cv::FILLED);
}
